from __future__ import annotations
import subprocess
from pathlib import Path
from typing import Generator
from datetime import datetime
from multiprocessing import Pool, Lock
import uuid

def get_time_string() -> str:
    return datetime.now().strftime('%Y/%m/%d, %H:%M:%S')

def log_status(dockin_file_name: str, status_code: int) -> None:
    lock.acquire()
    with open(root_path / 'log.csv', 'a+', encoding='utf-8') as f:
        f.write(f'{dockin_file_name},{status_code},{get_time_string()}\n')
    lock.release()

def pool_init(_lock):
    global lock
    lock = _lock

def get_unfinished_task(root_path: Path, result_root_path: Path) -> list[Path]:
    total_tasks = {}
    for dockin_file_path in root_path.glob('*.in'):
        total_tasks[dockin_file_path.stem] = dockin_file_path
    if result_root_path.exists():
        for protein_dir in result_root_path.iterdir():
            for binding_pocket_dir in protein_dir.iterdir():
                name = f'{protein_dir.stem}_{binding_pocket_dir.stem}'
                del total_tasks[name]
    return list(total_tasks.values())

def work(dockin_file_path: Path) -> None:
    uid = uuid.uuid4().hex
    with open(root_path / f'{uid}.in', 'w+', encoding='utf-8', newline='\n') as f:
        f.write(dockin_file_path.read_text().replace('ligands.list', f'{uid}.list'))
    with open(root_path / f'{uid}.list', 'w+', encoding='utf-8', newline='\n') as f:
        f.write((root_path/'ligands.list').read_text().replace('mols', uid))
    (root_path / uid).mkdir()
    for mol_file in (root_path/'mols').iterdir():
        (Path(uid) / mol_file.name).hardlink_to(mol_file)
    protein_name, docking_pocket_rank = dockin_file_path.stem.split('_')
    try:
        subprocess.run(['ledock', f'{uid}.in'])
        result_dir = result_root_path / protein_name / docking_pocket_rank
        result_dir.mkdir(parents=True, exist_ok=True)
        for dock_res in (root_path / uid).glob('*.dok'):
            dock_res.rename(result_dir/dock_res.name)
    except:
        log_status(dockin_file_path.stem, -1)
        return None
    finally:
        subprocess.run(['rm', f'{uid}.in'])
        subprocess.run(['rm', f'{uid}.list'])
        subprocess.run(['rm', '-r', uid])
    log_status(dockin_file_path.stem, 1)


if __name__ == '__main__':
    root_path = Path('.')
    result_root_path = Path('./result/')
    _lock = Lock()

    with Pool(initializer=pool_init, initargs=(_lock,)) as pool:
        pool.map(work, get_unfinished_task(root_path, result_root_path))